package com.sw.funfuck.common.util;

import com.sw.funfuck.common.util.ESQueryBuilder.DateHistogram;
import org.elasticsearch.action.ActionFuture;
import org.elasticsearch.action.bulk.BulkRequestBuilder;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.get.MultiGetItemResponse;
import org.elasticsearch.action.get.MultiGetResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.transport.TransportAddress;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.TermQueryBuilder;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.DeleteByQueryAction;
import org.elasticsearch.index.reindex.UpdateByQueryAction;
import org.elasticsearch.index.reindex.UpdateByQueryRequestBuilder;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.bucket.histogram.DateHistogramAggregationBuilder;
import org.elasticsearch.search.aggregations.bucket.histogram.ExtendedBounds;
import org.elasticsearch.search.aggregations.bucket.histogram.InternalDateHistogram;
import org.elasticsearch.search.aggregations.bucket.nested.Nested;
import org.elasticsearch.search.aggregations.bucket.nested.NestedAggregationBuilder;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.bucket.terms.Terms.Bucket;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.FieldSortBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.elasticsearch.transport.client.PreBuiltTransportClient;
import org.joda.time.DateTimeZone;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import java.net.InetAddress;
import java.util.*;

/**
 * @author chenxiaojian
 * @ClassName: ESUtil
 * @Description: Elasticsearch工具类
 * @date 2018年1月9日 下午5:03:34
 */
@Component
public class ESUtil {

    private static final Logger logger = LoggerFactory.getLogger(ESUtil.class);

    private final static int MAX = 10000;

    private TransportClient client;

    @Value("${es.cluster.name:elasticsearch}")
    private String clusterName;

    @Value("${es.cluster.address:127.0.0.1:9300}")
    private String clusterAddress;


    /**
     * client初始化
     */
    @PostConstruct
    public void init() {
        try {
            Settings settings = Settings.builder().put("cluster.name", clusterName).build();

            /**
             * 如果安装了x-pack组件，需要引入 x-pack-transport 使用PreBuiltXPackTransportClient实现 增加认证参数
             * https://www.elastic.co/guide/en/x-pack/current/java-clients.html
             */
            this.client = new PreBuiltTransportClient(settings);
            String[] nodes = clusterAddress.split(",");
            for (String node : nodes) {
                if (node.length() > 0) {
                    String[] hostPort = node.split(":");
                    this.client.addTransportAddress(
                            new TransportAddress(InetAddress.getByName(hostPort[0]), Integer.parseInt(hostPort[1])));
                }
            }
        } catch (Exception e) {
            logger.error("es init failed!", e);
        }
    }

    /**
     * 插入数据
     *
     * @param index  索引名
     * @param type   类型
     * @param id     数据id
     * @param source 数据
     */
    public void insertData(String index, String type, String id, Map<String, ?> source)
            throws Exception {
        client.prepareIndex(index, type).setId(id)
                .setSource(source)
                .get();
    }

    /**
     * 更新数据
     *
     * @param index 索引名
     * @param type  类型
     * @param id    数据id
     * @param json  数据
     */
    public void updateData(String index, String type, String id, String json)
            throws Exception {
        try {
            UpdateRequest updateRequest = new UpdateRequest(index, type, id)
                    .doc(json);
            client.update(updateRequest).get();
        } catch (Exception e) {
            throw new Exception("update data failed.", e);
        }
    }

    /**
     * 更新符合查询条件的数据
     *
     * @param index
     * @param queryBuilder
     * @param script
     * @return
     * @throws Exception
     */
    public long updateByQuery(String index, QueryBuilder queryBuilder, String script)
            throws Exception {
        try {
            UpdateByQueryRequestBuilder updateByQuery = UpdateByQueryAction.INSTANCE.newRequestBuilder(client);
            updateByQuery.source(index)
                    .filter(queryBuilder)
                    .size(1)
                    .script(new Script(ScriptType.INLINE, "painless", script, Collections.emptyMap()));
            BulkByScrollResponse response = updateByQuery.get();
            return response.getUpdated();
        } catch (Exception e) {
            throw new Exception("update data failed.", e);
        }
    }

    /**
     * 删除数据
     *
     * @param index 索引名
     * @param type  类型
     * @param id    数据id
     */
    public void deleteData(String index, String type, String id)
            throws Exception {
        client.prepareDelete(index, type, id)
                .get();
    }

    /**
     * 删除数据
     *
     * @param index   索引名
     * @param type    类型
     * @param builder ESQueryBuilder
     */
    public long deleteData(String index, String type, ESQueryBuilder builder)
            throws Exception {
        BulkByScrollResponse response = DeleteByQueryAction.INSTANCE.newRequestBuilder(client)
                .filter(builder.listBuilders()).source(index)
                .get();
        long deleted = response.getDeleted();
        return deleted;
    }

    /**
     * 根据批判属性删除
     * @param index
     * @param type
     * @param matchQueryBuilder
     * @return
     */
    public long deleteData(String index, String type, TermQueryBuilder matchQueryBuilder) {
        BulkByScrollResponse response = DeleteByQueryAction.INSTANCE.newRequestBuilder(client)
                .filter(matchQueryBuilder)
                .source(index)
                .get();
        return response.getDeleted();
    }


    /**
     * 批量插入数据
     *
     * @param index 索引名
     * @param type  类型
     * @param data  (_id 主键, json 数据)
     */
    public void bulkInsertData(String index, String type, Map<String, Object> data)
            throws Exception {
        BulkRequestBuilder bulkRequest = client.prepareBulk();
        data.forEach((param1, param2) -> {
            bulkRequest.add(client.prepareIndex(index, type, param1)
                    .setSource(param2)
            );
        });
        bulkRequest.get();
    }

    /**
     * 批量插入数据
     *
     * @param index 索引名
     * @param type  类型
     * @param data  批量数据
     */
    public void bulkInsertData(String index, String type, List<Map<String, Object>> data)
            throws Exception {
        BulkRequestBuilder bulkRequest = client.prepareBulk();
        data.forEach(item -> {
            bulkRequest.add(client.prepareIndex(index, type).setSource(item)
            );
        });
        bulkRequest.get();
    }

    /**
     * 查询
     *
     * @param index 索引名
     * @param type  类型
     * @param id    id
     */
    public String searchById(String index, String type, String id)
            throws Exception {
        return client.prepareGet(index, type, id).get().getSourceAsString();
    }

    /**
     * 查询
     *
     * @param index 索引名
     * @param type  类型
     * @param id    查询id
     */
    public Map<String, Object> searchMapById(String index, String type, String id)
            throws Exception {
        return client.prepareGet(index, type, id).get().getSourceAsMap();
    }

    /**
     * id查询
     *
     * @param index 索引名
     * @param type  类型
     * @param id    文档id
     * @return
     */
    public Map<String, Object> multiGetById(String index, String type, String id) {
        MultiGetResponse multiGetItemResponses = client.prepareMultiGet()
                .add(index, type, id)
                .get();

        for (MultiGetItemResponse itemResponse : multiGetItemResponses) {
            GetResponse response = itemResponse.getResponse();
            if (response.isExists()) {
                return response.getSourceAsMap();

            }
        }
        return null;
    }

    /**
     * 多id查询
     *
     * @param index 索引名
     * @param type  类型
     * @param ids   文档id集合
     * @return
     */
    public List<Map<String, Object>> multiGetByIds(String index, String type, List<String> ids) {
        MultiGetResponse multiGetItemResponses = client.prepareMultiGet()
                .add(index, type, ids)
                .get();
        List<Map<String, Object>> mapList = new ArrayList<>();
        for (MultiGetItemResponse itemResponse : multiGetItemResponses) {
            GetResponse response = itemResponse.getResponse();
            if (response.isExists()) {
                mapList.add(response.getSourceAsMap());

            }
        }
        return mapList;
    }

    /**
     * 查询
     *
     * @param index 索引名
     * @param type  类型
     * @param from  from
     * @param size  size
     */
    public SearchResponse search(String index, String type, int from, int size, QueryBuilder queryBuilder) {
        SearchResponse searchResponse = client.prepareSearch(index).setTypes(type).setQuery(queryBuilder).setFrom(from).setSize(size).get();
        return searchResponse;
    }

    /**
     * 查询
     *
     * @param index               索引名
     * @param type                类型
     * @param searchSourceBuilder 查询构造器
     */
    public SearchResponse search(String index, String type, SearchSourceBuilder searchSourceBuilder) {
//        BoolQueryBuilder boolBuilder = new BoolQueryBuilder();
//        MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery("detectTasks.detectType","ALARM");
//        NestedQueryBuilder nestedQuery = new NestedQueryBuilder("",matchQueryBuilder, ScoreMode.Avg);
//        boolBuilder.must(nestedQuery);
//
//        AggregationBuilder sumAgg = AggregationBuilders
//                .nested("sumnest","detectTasks.similars")
//                .subAggregation(AggregationBuilders
//                        .sum("sumagg")
//                        .field("target"));
//
//        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
//        searchSourceBuilder
//                .size(10)
//                .from(10)
//                .query(boolBuilder)
//                .aggregation(sumAgg);

        //searchSourceBuilder.toString();

        SearchRequest searchRequest = new SearchRequest();
        searchRequest.indices("index");
        searchRequest.types("_doc");
        searchRequest.source(searchSourceBuilder);

        //searchRequest.toString();

        ActionFuture<SearchResponse> actionFuture = this.client.search(searchRequest);
        return actionFuture.actionGet();
    }

    /**
     * 查询
     *
     * @param index   索引名
     * @param type    类型
     * @param builder 查询构造
     */
    public List<String> search(String index, String type, ESQueryBuilder builder)
            throws Exception {
        List<String> result = new ArrayList<>();
        SearchRequestBuilder searchRequestBuilder = createSearchRequestBuilder(index, type, builder)
                .addSort(createSortBuilder(builder));
        SearchResponse sr = searchRequestBuilder.execute().actionGet();

        SearchHits hits = sr.getHits();
        SearchHit[] searchHists = hits.getHits();
        for (SearchHit sh : searchHists) {
            result.add(sh.getSourceAsString());
        }
        return result;
    }

    /**
     * 查询
     *
     * @param index   索引名
     * @param type    类型
     * @param builder 查询构造
     */
    public SearchResponse searchResponse(String index, String type, ESQueryBuilder builder)
            throws Exception {
        SearchRequestBuilder searchRequestBuilder = createSearchRequestBuilder(index, type, builder)
                .addSort(createSortBuilder(builder));
        SearchResponse sr = searchRequestBuilder.execute().actionGet();

        return sr;
    }

    /**
     * 查询
     *
     * @param index        索引名
     * @param type         类型
     * @param queryBuilder 查询构造
     */
    public SearchResponse searchResponse(String index, String type, QueryBuilder queryBuilder, FieldSortBuilder sortBuilder)
            throws Exception {
        SearchRequestBuilder searchRequestBuilder = client.prepareSearch(index).setTypes(type).setQuery(queryBuilder);
        if (sortBuilder != null) {
            searchRequestBuilder.addSort(sortBuilder);
        }

        SearchResponse sr = searchRequestBuilder.execute().actionGet();
        return sr;
    }

    /**
     * 查询
     *
     * @param index        索引名
     * @param type         类型
     * @param queryBuilder 查询构造
     */
    public SearchResponse searchResponse(String index, String type, QueryBuilder queryBuilder, FieldSortBuilder sortBuilder, Integer from, Integer size)
            throws Exception {
        SearchRequestBuilder searchRequestBuilder = client.prepareSearch(index).setTypes(type).setQuery(queryBuilder);
        if (sortBuilder != null) {
            searchRequestBuilder.addSort(sortBuilder);
        }
        if (size != null) {
            size = size < 0 ? 0 : size;
            size = size > MAX ? MAX : size;
            //返回条目数 
            searchRequestBuilder.setSize(size);
        }
        if (from != null) {
            searchRequestBuilder.setFrom(from);
        }
        SearchResponse sr = searchRequestBuilder.execute().actionGet();
        return sr;
    }


    /**
     * 查询
     *
     * @param index   索引名
     * @param type    类型
     * @param builder 查询构造
     */
    public List<String> searchStr(String index, String type, QueryBuilder builder, FieldSortBuilder sortBuilder, Integer from, Integer size)
            throws Exception {
        List<String> ret = new LinkedList<>();
        SearchResponse sr = searchResponse(index, type, builder, sortBuilder, from, size);
        SearchHits hits = sr.getHits();
        SearchHit[] searchHists = hits.getHits();
        for (SearchHit sh : searchHists) {
            ret.add(sh.getSourceAsString());
        }
        return ret;
    }

    /**
     * 查询单条记录
     *
     * @param index
     * @param type
     * @param builder
     * @return
     */
    public String searchOne(String index, String type, ESQueryBuilder builder) throws Exception {
        SearchRequestBuilder searchRequestBuilder = createSearchRequestBuilder(index, type, builder)
                .addSort(createSortBuilder(builder));
        searchRequestBuilder.setSize(1);
        SearchResponse searchResponse = searchRequestBuilder.get();
        SearchHits hits = searchResponse.getHits();
        SearchHit[] searchHists = hits.getHits();
        if (searchHists.length == 0) {
            return null;
        }
        return searchHists[0].getSourceAsString();
    }

    /**
     * 计数
     *
     * @param index   索引名
     * @param type    类型
     * @param builder 查询构造
     */
    public long statCount(String index, String type, ESQueryBuilder builder) throws Exception {
        SearchRequestBuilder searchRequestBuilder = client.prepareSearch(index).setTypes(type)
                .setQuery(builder.listBuilders());
        SearchResponse searchResponse = searchRequestBuilder.execute().actionGet();
        return searchResponse.getHits().totalHits;
    }

    /**
     * 计数
     *
     * @param index   索引名
     * @param type    类型
     * @param builder 查询构造
     */
    public long statCount(String index, String type, QueryBuilder builder) throws Exception {
        SearchRequestBuilder searchRequestBuilder = client.prepareSearch(index).setTypes(type)
                .setQuery(builder)
                .setSize(0);
        SearchResponse searchResponse = searchRequestBuilder.execute().actionGet();
        return searchResponse.getHits().totalHits;
    }

    /**
     * 分组统计查询(只统计每个分组数量, 不包含每个分组下的数据项)
     *
     * @param index
     * @param type
     * @param builder
     * @param groupBy
     * @return
     */
    @SuppressWarnings("unchecked")
    public Map<Object, Object> statSearch(String index, String type, ESQueryBuilder builder, String groupBy) throws Exception {

        Map<Object, Object> map = new HashMap<>();
        builder = new ESQueryBuilder();
        builder.setQueryBuilder(map);
        SearchRequestBuilder searchRequestBuilder = createSearchRequestBuilder(index, type, builder)
                .addSort(createSortBuilder(builder));
        searchRequestBuilder.addAggregation(
                AggregationBuilders.terms("agg").field(groupBy)
        );
        SearchResponse sr = searchRequestBuilder.get();
        Terms agg = sr.getAggregations().get("agg");
        Iterator<Bucket> iter = (Iterator<Bucket>) agg.getBuckets().iterator();
        while (iter.hasNext()) {
            Bucket gradeBucket = iter.next();
            map.put(gradeBucket.getKey(), gradeBucket.getDocCount());
        }
        return map;
    }

    public List<InternalDateHistogram.Bucket> statSearchByDateHistogram(String index, String type, QueryBuilder builder, DateHistogram dh)
            throws Exception {


        SearchRequestBuilder searchRequestBuilder = client.prepareSearch(index).setTypes(type);
        searchRequestBuilder.setQuery(builder);
        DateHistogramAggregationBuilder dhab = AggregationBuilders
                .dateHistogram("agg")
                .field(dh.getFieldName())
                .dateHistogramInterval(dh.getDateInterval())
                .format(dh.getFormat())
                .timeZone(DateTimeZone.forID("Asia/Shanghai"));
        dhab.minDocCount(0);
        if (dh.getMin() != null && dh.getMax() != null) {
            dhab.extendedBounds(new ExtendedBounds(dh.getMin(), dh.getMax()));
        }
        searchRequestBuilder.addAggregation(dhab);
        searchRequestBuilder.setSize(0);

        SearchResponse sr = searchRequestBuilder.get();
        List<InternalDateHistogram.Bucket> buckets = ((InternalDateHistogram) sr.getAggregations().get("agg")).getBuckets();
        return buckets;
    }

    public List<InternalDateHistogram.Bucket> statSearchByNestDateHistogram(String index, String type, QueryBuilder builder, DateHistogram dh, String nestPath)
            throws Exception {

        NestedAggregationBuilder nestedQueryBuilder = AggregationBuilders.nested("negsted", nestPath);

        SearchRequestBuilder searchRequestBuilder = client.prepareSearch(index).setTypes(type);
        searchRequestBuilder.setQuery(builder);
        DateHistogramAggregationBuilder dhab = AggregationBuilders
                .dateHistogram("agg")
                .field(dh.getFieldName())
                .dateHistogramInterval(dh.getDateInterval())
                .format(dh.getFormat())
                .timeZone(DateTimeZone.forID("Asia/Shanghai"));
        dhab.minDocCount(0);
        if (dh.getMin() != null && dh.getMax() != null) {
            dhab.extendedBounds(new ExtendedBounds(dh.getMin(), dh.getMax()));
        }

        nestedQueryBuilder.subAggregation(dhab);
        searchRequestBuilder.addAggregation(nestedQueryBuilder);
        searchRequestBuilder.setSize(0);

        SearchResponse sr = searchRequestBuilder.get();

        //获取一级嵌套
        Nested ntd = sr.getAggregations().get("negsted");

        //从一级嵌套查询的结果中获取二级嵌套查询结果
        InternalDateHistogram term = ntd.getAggregations().get("agg");

        List<InternalDateHistogram.Bucket> buckets = term.getBuckets();
        return buckets;
    }


    /**
     * 构造QueryBuilder
     *
     * @param index
     * @param type
     * @param builder
     * @return
     */
    private SearchRequestBuilder createSearchRequestBuilder(String index, String type, ESQueryBuilder builder) throws Exception {
        SearchRequestBuilder searchRequestBuilder = client.prepareSearch(index).setTypes(type);
        searchRequestBuilder.setQuery(builder.listBuilders());
        int size = builder.getSize();
        if (size < 0) {
            size = 0;
        }
        if (size > MAX) {
            size = MAX;
        }
        //返回条目数  
        builder.setSize(size);
        builder.setFrom(builder.getFrom() < 0 ? 0 : builder.getFrom());
        searchRequestBuilder.setSize(size);
        searchRequestBuilder.setFrom(builder.getFrom());
        return searchRequestBuilder;
    }


    /**
     * 构造排序Builder
     *
     * @param queryBuilder
     * @return
     */
    private FieldSortBuilder createSortBuilder(ESQueryBuilder queryBuilder) {
        FieldSortBuilder sortBuilder = null;
        if (queryBuilder.getAsc() != null && queryBuilder.getAsc().length() > 0) {
            sortBuilder = new FieldSortBuilder(queryBuilder.getAsc());
            sortBuilder.order(SortOrder.ASC);
            sortBuilder.unmappedType("date");
        }
        if (queryBuilder.getDesc() != null && queryBuilder.getDesc().length() > 0) {
            sortBuilder = new FieldSortBuilder(queryBuilder.getDesc());
            sortBuilder.order(SortOrder.DESC);
            sortBuilder.unmappedType("date");
        }
        return sortBuilder;
    }

    /**
     * 关闭链接
     */
    public void close() {
        client.close();
    }

    /**
     *
     * @ClassName: ESQuery
     * @Description: 配置ES查询注解
     * @author chenxiaojian
     * @date 2018年1月9日 下午4:09:30
     *
     */
    @Retention(value = RetentionPolicy.RUNTIME)
    @Target(value = ElementType.FIELD)
    public static @interface ESQuery {

        /**
         * 配置ES document属性名
         * 不配置则默认字段名
         */
        public String key() default "";

        /**
         * 包括：
         *  term  默认精确匹配字段
         * 	terms 多值查询
         * 	rangeStart, rangeEnd 范围查询
         *  match 分词查询
         */
        public ESQueryBuilder.ESQueryType type() default ESQueryBuilder.ESQueryType.TERM;


        /**
         * type为terms和range时, 对值进行拆分
         * 默认,
         */
        public String split() default ",";

    }
}